#include "v1model.p4"

#define MAX_PORTS 256 // 交换机端口的数量
#define CPU_PORT 510 // cpu端口号
#define MAX_HOPS 5 // 路由跳数限制

// 类型常量定义
const bit<16> TYPE_IPV4 = 0x0800;
const bit<16> TYPE_SPROBE = 0x0810;
const bit<8> TYPE_TCP =0x06;
const bit<8> TYPE_UDP =0x11;

typedef bit<9>  egressSpec_t;
typedef bit<4> ce_t;
typedef bit<8> route_index_t;
typedef bit<16> ip_index_t;
typedef bit<2> rv_t;


// 寄存器容量定义
const bit<32> NUMBER_65536 = 0xFFFF;
const bit<32> NUMBER_256 = MAX_PORTS;
const bit<32> NUMBER_2024 = 2024;

typedef bit<16> PortIdToController_t;

/*************************************************************************
*********************** H E A D E R S  ***********************************
*************************************************************************/

// 以太网报头
header ethernet_h {
    bit<48> dstAddr;
    bit<48> srcAddr;
    bit<16> etherType;
}

// IPv4报头
header ipv4_h {
    bit<4>  version;
    bit<4>  ihl;
    bit<8>  diffserv;  // 用于存储route_index
    bit<16> totalLen;
    bit<16> identification;
    bit<3>  flags;
    bit<13> fragOffset;
    bit<8>  ttl;
    bit<8>  protocol;
    bit<16> hdrChecksum;
    bit<32> srcAddr;
    bit<32> dstAddr;
}

// S_Probe探针报头
header s_probe_h {
    bit<32> src_ip;
    bit<32> dst_ip;
    bit<8> route_index;
    bit<4>  direction; // 0:源->目的, 1:目的->源
    bit<4> has_metadata;
}

// 元数据段
header metadata_h {
    bit<32> enq_qdepth;  // 队列长度
    bit<32> deq_qdepth;
    bit<32> enq_timestamp;
    bit<32> deq_timestamp;
    bit<8> next;
}

header tcp_h {
    bit<16> srcPort;
    bit<16> dstPort;
    bit<32> seqNo;
    bit<32> ackNo;
    bit<4>  dataOffset;
    bit<3>  res;
    bit<3>  ecn;
    bit<6>  ctrl;
    bit<16> window;
    bit<16> checksum;
    bit<16> urgentPtr;
}

header udp_h {
    bit<16> src_port;  // Source port
    bit<16> dst_port;  // Destination port
    bit<16> length;    // Length of UDP packet
    bit<16> checksum;  // Checksum
}

// 和控制层进行交互的信息格式
enum bit<8> ControllerOpcode_e {
    NO_OP          = 0,
    READ_CONGESTION_REG  = 1,
    WRITE_CONGESTION_REG = 2,
    READ_FLOWTABLE_REG =3,
    RESET_FLOWTABLE_REG = 4,
    WRITE_FLOWTABLE_REG = 5,
    GENERATE_SPROBE = 6,
    UPDATE_CONGESTION_MSG = 7
}

enum bit<8> PuntReason_e {
    UNRECOGNIZED_OPCODE = 1,
    OPERATION_RESPONSE  = 2,
    ACTIVE_REQUEST = 3,
}

@controller_header("packet_in")
header packet_in_header_h {
    PortIdToController_t input_port;
    PuntReason_e         punt_reason;
    ControllerOpcode_e   opcode;
    bit<32> operand0;
    bit<32> operand1;
    bit<32> operand2;
    bit<32> operand3;
    bit<32> operand4;
    bit<32> operand5;
    bit<32> operand6;
    bit<32> operand7;
    bit<32> operand8;
    bit<32> operand9;
    bit<32> operand10;
    bit<32> operand11;
    bit<32> operand12;
    bit<32> operand13;
    bit<32> operand14;
    bit<32> operand15;
    bit<32> operand16;
    bit<32> operand17;
    bit<32> operand18;
}

@controller_header("packet_out")
header packet_out_header_h {
    PortIdToController_t input_port;
    ControllerOpcode_e   opcode;
    bit<32> operand0;
    bit<32> operand1;
    bit<32> operand2;
    bit<32> operand3;
    bit<32> operand4;
    bit<32> operand5;
    bit<32> operand6;
}

struct metadata{
    ip_index_t ip_index;
    bit <16> hash_value;
    route_index_t route_index;
    bit<16> port; // TCP或者UDP的端口号
    bit<32> port_number; // 端口已经转发的字节数量
    bit <4> random_value;
    bit <32> ipv4; // 保存IPv4地址
}

struct headers {
    packet_in_header_h packet_in;
    packet_out_header_h packet_out;
    ethernet_h ethernet;
    s_probe_h probe;
    metadata_h [MAX_HOPS] metadatas;
    ipv4_h ipv4;
    tcp_h tcp;
    udp_h udp;
}


/*************************************************************************
*********************** P A R S E R  ***********************************
*************************************************************************/

parser MyParser(packet_in packet,
                out headers hdr,
                inout metadata meta,
                inout standard_metadata_t standard_metadata) {
    
    state start {
        transition parse_packet;
    }

    state parse_packet{

        transition select(standard_metadata.ingress_port){
            // modify: 
            CPU_PORT: parse_packet_out_of_controller;
            default: parse_ethernet;
        }
    }

    state parse_ethernet{
        
        packet.extract(hdr.ethernet);
        transition select(hdr.ethernet.etherType){

            TYPE_IPV4: parse_ipv4;
            TYPE_SPROBE: parse_s_probe;
            default: accept;
        }
    }

    state parse_s_probe{
      packet.extract(hdr.probe);
      transition select(hdr.probe.has_metadata){
        1: parse_metadatas;
        0:accept;
      }
    }

    state parse_metadatas{
      packet.extract(hdr.metadatas.next);
      transition select (hdr.metadatas.last.next){
        1: parse_metadatas;
        0: accept;
        default: accept;
      }
    }

    state parse_packet_out_of_controller{
        packet.extract(hdr.packet_out);
        transition accept;
    }

    state parse_ipv4{
        packet.extract(hdr.ipv4);
        transition select(hdr.ipv4.protocol){
            TYPE_TCP: parse_tcp;
            TYPE_UDP: parse_udp;
            default: accept;
        }
    }
    state parse_tcp{
        packet.extract(hdr.tcp);
        transition accept;
    }

    state parse_udp{
        packet.extract(hdr.udp);
        transition accept;
    }
}

/*************************************************************************
************   C H E C K S U M    V E R I F I C A T I O N   *************
*************************************************************************/

control MyVerifyChecksum(inout headers hdr, inout metadata meta) {
    apply { 
        
    }
}

/*************************************************************************
**************  I N G R E S S   P R O C E S S I N G   *******************
*************************************************************************/

control MyIngress(inout headers hdr,
                  inout metadata meta,
                  inout standard_metadata_t standard_metadata) {
    // 寄存器定义
    register<route_index_t>(NUMBER_256) main_route_register; // 主路由表
    register<route_index_t>(NUMBER_256) slave_route_register; // 从路由表
    register<bit<32>>(NUMBER_256) output_port_bytes_register; // 端口字节计数器
    register<ce_t>(NUMBER_65536) congestion_register;        // 拥塞寄存器
    register<egressSpec_t>(NUMBER_2024) route_index_to_port_register;

    register<bit<32>>(4) metadata_register; // 保存参数数据 0 -> 端口转发上限值; 1-> 拥塞值上限; 2 -> 拥塞值上限; 3 -> 端口缩小因子
    // 路由表定义
    action set_ip_index(bit<16> ip_index) {
      meta.ip_index = ip_index;
    }
    table ip_to_index_table {
        key = {
            hdr.ipv4.dstAddr: lpm;
            meta.port: exact;
        }
        actions = { 
            set_ip_index; 
            NoAction;
        }
        default_action = NoAction;
        size = 1024;
    }
    action ipv4_forward(bit<8> port) { 
        standard_metadata.egress_spec = (egressSpec_t)port;
        hdr.ipv4.ttl = hdr.ipv4.ttl - 1;
    }
    table ipv4_forward_table {
        key = { 
            hdr.ipv4.dstAddr: lpm;
        }
        actions = {
            ipv4_forward;
            NoAction;
        }
        default_action = NoAction;
        size = 1024;
    }

    action set_mac(bit<48> mac_dst) {
        hdr.ethernet.srcAddr = hdr.ethernet.dstAddr;
        hdr.ethernet.dstAddr = mac_dst;
    }
    table mac_table {
        key = { 
            standard_metadata.egress_spec: exact; 
        }
        actions = {
            set_mac;
            NoAction;
        }
        default_action = NoAction;
        size = 256;
    }

    action is_dst_switch(){
        // probe报文到达目的交换机之后，将probe报文发送给控制器
        standard_metadata.egress_spec = CPU_PORT;

        // 报文封装
        hdr.packet_in.setValid();
        hdr.packet_in.input_port = (PortIdToController_t)standard_metadata.egress_spec;
        hdr.packet_in.punt_reason = PuntReason_e.ACTIVE_REQUEST;
        hdr.packet_in.opcode = ControllerOpcode_e.UPDATE_CONGESTION_MSG;
        hdr.packet_in.operand0 = hdr.probe.src_ip;
        hdr.packet_in.operand1 = hdr.probe.dst_ip;
        hdr.packet_in.operand2 = (bit<32>)hdr.probe.route_index;

        
        hdr.packet_in.operand3 = hdr.metadatas[0].enq_timestamp;
        hdr.packet_in.operand4 = hdr.metadatas[0].deq_timestamp;
        hdr.packet_in.operand5 = hdr.metadatas[0].enq_qdepth;
        hdr.packet_in.operand6 = hdr.metadatas[0].deq_qdepth;

        hdr.packet_in.operand7 = hdr.metadatas[1].enq_timestamp;
        hdr.packet_in.operand8 = hdr.metadatas[1].deq_timestamp;
        hdr.packet_in.operand9 = hdr.metadatas[1].enq_qdepth;
        hdr.packet_in.operand10 = hdr.metadatas[1].deq_qdepth;

        hdr.packet_in.operand11 = hdr.metadatas[2].enq_timestamp;
        hdr.packet_in.operand12 = hdr.metadatas[2].deq_timestamp;
        hdr.packet_in.operand13 = hdr.metadatas[2].enq_qdepth;
        hdr.packet_in.operand14 = hdr.metadatas[2].deq_qdepth;

        hdr.packet_in.operand15 = hdr.metadatas[3].enq_timestamp;
        hdr.packet_in.operand16 = hdr.metadatas[3].deq_timestamp;
        hdr.packet_in.operand17 = hdr.metadatas[3].enq_qdepth;
        hdr.packet_in.operand18 = hdr.metadatas[3].deq_qdepth;

                // TODO: 封装metadata
        hdr.metadatas[0].setInvalid();
        hdr.metadatas[1].setInvalid();
        hdr.metadatas[2].setInvalid();
        hdr.metadatas[3].setInvalid();

        hdr.ethernet.setInvalid();
        hdr.probe.setInvalid();

    }

    table switch_type_table{
        key={
            hdr.probe.dst_ip: lpm;
        }
        actions = {
            is_dst_switch;
            NoAction;
        }
        default_action = NoAction;
        size = 256;
    }

    // 由route_index 映射成端口号port
    action route_to_index_port(){
        route_index_to_port_register.read(standard_metadata.egress_spec,(bit<32>)meta.route_index);
    }

    // 计算备选路径
    action compute_slave_route_index(){
        bit<32> congestion_threshold;
        metadata_register.read(congestion_threshold,(bit<32>)2);

        bit<32> random_value_threshold;
        metadata_register.read(random_value_threshold,(bit<32>)1);

        rv_t random_value = (rv_t)(meta.hash_value);
        slave_route_register.read(meta.route_index,(bit<32>)(meta.ip_index++random_value));
        ce_t ce1;
        congestion_register.read(ce1,(bit<32>)(meta.ip_index++meta.route_index));
        // 进行两次路由探测
        if(ce1 > (ce_t)congestion_threshold){
            random_value = random_value+1;
            if(random_value == (rv_t)random_value_threshold){
                random_value = 0;
            }
            slave_route_register.read(meta.route_index,(bit<32>)(meta.ip_index++random_value));
            congestion_register.read(ce1,(bit<32>)(meta.ip_index++meta.route_index));
            if(ce1 > (ce_t)congestion_threshold){
                random_value = random_value+1;
                if(random_value == (rv_t)random_value_threshold){
                    random_value = 0;
                }
                slave_route_register.read(meta.route_index,(bit<32>)(meta.ip_index++random_value));
            }
        }
    }

    apply {
        if (hdr.ipv4.isValid()) {
            meta.ipv4 = hdr.ipv4.dstAddr;
            if (ipv4_forward_table.apply().miss) {
                // 检查三层端口地址
                if(hdr.udp.isValid()){
                    meta.port  = hdr.udp.dst_port;
                }else if(hdr.tcp.isValid()){
                    meta.port = hdr.tcp.dstPort;
                }else{
                    meta.port = 0; // 没有采用传输层协议时端口号为0
                }
                ip_to_index_table.apply(); // 获取IP_INDEX
                bit<16> hash_value;
                if(hdr.udp.isValid()){
                  hash(hash_value,
                    HashAlgorithm.crc16,
                    (bit<16>)0,
                    {hdr.ipv4.srcAddr,hdr.ipv4.dstAddr,hdr.udp.src_port,hdr.udp.dst_port},
                    (bit<16>)65535);
                }else if(hdr.tcp.isValid()){
                  hash(hash_value,
                    HashAlgorithm.crc16,
                    (bit<16>)0,
                    {hdr.ipv4.srcAddr,hdr.ipv4.dstAddr,hdr.tcp.srcPort,hdr.tcp.dstPort},
                    (bit<16>)65535);
                }else{
                  hash(hash_value,
                    HashAlgorithm.crc16,
                    (bit<16>)0,
                    {hdr.ipv4.srcAddr,hdr.ipv4.dstAddr,hdr.ipv4.protocol},
                    (bit<16>)65535);
                }
                meta.hash_value = hash_value;
                route_index_t route_index;
                main_route_register.read(route_index,(bit<32>)meta.ip_index);
                // 获取输出端口号
                meta.route_index = route_index;
                route_to_index_port();
                // 检查端口转发字节数量
                output_port_bytes_register.read(meta.port_number,(bit<32>)standard_metadata.egress_spec);
                egressSpec_t output_number  = standard_metadata.egress_spec;

                bit<32> port_forward_bytes_threshold;
                metadata_register.read(port_forward_bytes_threshold,(bit<32>)0);
                if(meta.port_number > port_forward_bytes_threshold){
                    bit<32> congestion_rv_threshold;
                    metadata_register.read(congestion_rv_threshold,(bit<32>)1);
                    compute_slave_route_index();// 计算备选路由
                    route_to_index_port(); // 得到输出端口
                    // 更新输出端口字节计数值
                    bit <32> sub_factor;
                    metadata_register.read(sub_factor,(bit<32>)3);
                    // 从路由表每转发一定字节数量数据包，更新主路由端口计数值
                    meta.port_number  = meta.port_number - sub_factor;
                    if(meta.port_number < 0){
                      meta.port_number = 0;
                    }

                }
                // 更新端口转发的字节数量
                bit<32> original_number;
                output_port_bytes_register.read(original_number,(bit<32>)output_number);
                output_port_bytes_register.write((bit<32>)output_number,meta.port_number+original_number);
                log_msg("sending ipv4 packet to port {}",{standard_metadata.egress_spec});
            }
            
        } else if (hdr.probe.isValid()) {
            // 探针处理逻辑
                meta.route_index = hdr.probe.route_index; // 根据路径标识进行转发
                route_to_index_port();
                // TODO: 添加元数据段，完成赋值操作
                hdr.metadatas.push_front(1);

                // TODO: 更新metadata报文
                hdr.metadatas[0].enq_qdepth = (bit<32>)standard_metadata.enq_qdepth;
                hdr.metadatas[0].deq_qdepth = (bit<32>)standard_metadata.deq_qdepth;
                hdr.metadatas[0].enq_timestamp = (bit<32>)standard_metadata.enq_timestamp;
                hdr.metadatas[0].deq_timestamp = (bit<32>)standard_metadata.deq_timedelta;
                hdr.metadatas[0].next  = 1;

                // TODO: 到达目的交换机之后将metadata信息发往控制器
                switch_type_table.apply();
        }else if(hdr.packet_out.isValid()){
          switch(hdr.packet_out.opcode){
            ControllerOpcode_e.GENERATE_SPROBE:{
              log_msg("receiving packets from controller to generate S_Probe packet");
              // 完成报文封装
              hdr.ethernet.setValid();
              hdr.ethernet.etherType = TYPE_SPROBE;
              hdr.probe.setValid();
              hdr.probe.src_ip = (bit<32>)hdr.packet_out.operand0;
              hdr.probe.dst_ip = (bit<32>)hdr.packet_out.operand1;
              hdr.probe.route_index = (bit<8>)hdr.packet_out.operand2;
              hdr.probe.direction = (bit<4>)hdr.packet_out.operand3;
              hdr.probe.has_metadata = (bit<4>)hdr.packet_out.operand4;
              hdr.packet_out.setInvalid();
              // 将报文从端口发出
              meta.route_index = hdr.probe.route_index; // 根据路径标识进行转发
              route_to_index_port();
              log_msg("sending probe src = {} ,dst = {},route_index = {} at port {}",
              {hdr.probe.src_ip,hdr.probe.dst_ip,hdr.probe.route_index,standard_metadata.egress_spec});
              // hdr.metadatas.setValid();
            }
            ControllerOpcode_e.UPDATE_CONGESTION_MSG:{
              log_msg("receiving packets from controller to update congestion message");
              // TODO: 更新congestion_register寄存器
            }
          }
        }
        // 设置mac地址
        mac_table.apply();
    }
}

/*************************************************************************
****************  E G R E S S   P R O C E S S I N G   *******************
*************************************************************************/

control MyEgress(inout headers hdr,
                 inout metadata meta,
                 inout standard_metadata_t standard_metadata) {
    apply {
        
    }
}

/*************************************************************************
*************   C H E C K S U M    C O M P U T A T I O N   **************
*************************************************************************/

control MyComputeChecksum(inout headers hdr, inout metadata meta) {
     apply {
        update_checksum(
            hdr.ipv4.isValid(),
            { hdr.ipv4.version,
              hdr.ipv4.ihl,
              hdr.ipv4.diffserv,
              hdr.ipv4.totalLen,
              hdr.ipv4.identification,
              hdr.ipv4.flags,
              hdr.ipv4.fragOffset,
              hdr.ipv4.ttl,
              hdr.ipv4.protocol,
              hdr.ipv4.srcAddr,
              hdr.ipv4.dstAddr },
            hdr.ipv4.hdrChecksum,
            HashAlgorithm.csum16);
    
    }
}


/*************************************************************************
***********************  D E P A R S E R  *******************************
*************************************************************************/

control MyDeparser(packet_out packet, in headers hdr) {

    apply{
        packet.emit(hdr.packet_in);
        packet.emit(hdr.ethernet);
        packet.emit(hdr.probe);
        packet.emit(hdr.metadatas);
        packet.emit(hdr.ipv4);
        packet.emit(hdr.tcp);
        packet.emit(hdr.udp);
    }
}

/*************************************************************************
***********************  S W I T C H  *******************************
*************************************************************************/

V1Switch(
MyParser(),
MyVerifyChecksum(),
MyIngress(),
MyEgress(),
MyComputeChecksum(),
MyDeparser()
) main;